package scu.maqiang.possion;

import scu.maqiang.fes.BVPType;
import scu.maqiang.fes.FES2T31;
import scu.maqiang.mesh.MatrixFunc;
import scu.maqiang.mesh.Mesh2T3;
import scu.maqiang.mesh.ScalarFunc;
import scu.maqiang.mesh.T3Type;
import scu.maqiang.numeric.Direct;
import scu.maqiang.numeric.DirectSSolver;
import scu.maqiang.numeric.MVO;
import scu.maqiang.numeric.SRMatrix;

/**
 * 二维线性三角形单元泊松问题求解器
 * 
 * @author 马强
 *
 */
public class Possion2DT31 {

	/**
	 * 泊松问题求解类
	 * @param mesh
	 */
	public Possion2DT31(Mesh2T3 mesh) {
		this.mesh = mesh;
		fs = new FES2T31(mesh);
		A = new SRMatrix(fs.GetNdof());
		x = new double[fs.GetNdof()];
		gradient = new double[2][fs.GetNdof()];
		RHS = new double[fs.GetNdof()];
	}

	/**
	 * 设置Possion问题网格
	 * @param mesh 二维线性三角形网格对象
	 */
	public void setMesh(Mesh2T3 mesh) {
		this.mesh = mesh;
		fs = new FES2T31(mesh);
		A = new SRMatrix(fs.GetNdof());
		x = new double[fs.GetNdof()];
		gradient = new double[2][fs.GetNdof()];
		RHS = new double[fs.GetNdof()];
	}

	/**
	 * 设置温度刚度矩阵, 热传导系数为常数
	 * @param coef 材料热传导系数
	 * @param type 问题类型, 选择为COMMON(直角坐标问题), AXISYMMETRIC(轴对称问题),
	 */
	public void setHeatStiff(double coef, BVPType type) {
		A.clear();
		fs.assembleStiff(new double[]{coef}, type, A);
	}

	/**
	 * 设置温度刚度矩阵, 热传导系数为二阶张量, 用一个二维数组表示
	 * @param tensor 二阶热传导张量，2x2数组
	 * @param type 问题类型, 选择为CONSTITUTE_MATRIX_COMMON(直角坐标问题), CONSTITUTE_MATRIX_AXIS(轴对称问题)
	 */
	public void setHeatStiff(double[][] tensor, BVPType type) {
		A.clear();
		fs.assembleStiff(tensor, type, A);
	}

	/**
	 * 设置温度刚度矩阵, 热传导系数为标量函数
	 * @param coefFunc 热传导系数函数接口, 函数可以使用lambda函数组出, 带传入的三个参数, 可定义如下
	 *                 ScalarFunc coefFunc = (xy, label, param) -> ......;
	 *                 其中xy表示网格结点坐标, label表示单元编号, param 为外设额外参数, 返回一个标量值
	 * @param param 外设额外参数
	 * @param type 问题类型, 选择为COMMON(直角坐标问题), AXISYMMETRIC(轴对称问题)
	 */
	public void setHeatStiff(ScalarFunc coefFunc, double[] param, BVPType type) {
		A.clear();
		fs.assembleStiff(new ScalarFunc[]{coefFunc}, param, type, A);
	}

	/**
	 * 设置温度刚度矩阵, 热传导系数为二阶矩阵函数, 用一个二维数组表示
	 * @param matCoef 二阶热传导系数矩阵，返回一个2x2数组, 带传入三个参数, 可定义如下:
	 *                  MatrixFunc matFunc = (xy, label, param) -> ......;
	 *                 其中xy表示网格结点坐标, label表示单元编号, param 为外设额外参数, 返回2x2数组
	 * @param type 问题类型, 选择为CONSTITUTE_MATRIX_COMMON(直角坐标问题), CONSTITUTE_MATRIX_AXIS(轴对称问题)
	 */
	public void setHeatStiff(MatrixFunc matCoef, BVPType type, double[][][] param) {
		A.clear();
		fs.assembleStiff(matCoef, param, type, A);
	}

	/**
	 * 施加常热源
	 * @param coef 热源常数
	 * @param type 问题类型, 与热刚度矩阵的问题类型保持一致
	 */
	public void applyHeatSource(double coef, BVPType type) {
		fs.assembleSource(new double[]{coef}, type, RHS);
	}

	/**
	 * 施加热源向量
	 * @param vecCoef 数组, 长度等于网格结点个数(问题自由度数)
	 * @param type 问题类型, 与热刚度矩阵的问题类型保持一致
	 */
	public void applyHeatSource(double[] vecCoef, BVPType type) {
		fs.assembleSource(new double[][]{vecCoef}, type, RHS);
	}


	/**
	 * 施加热源函数
	 * @param sourceFunc 热源函数接口, 函数可以使用lambda函数组出, 带传入的三个参数, 可定义如下
	 * 	 *               ScalarFunc sourceFunc = (xy, label, param) -> ......;
	 * 	 *               其中xy表示网格结点坐标, label表示单元编号, param 为外设额外参数, 返回一个标量值
	 * @param type 问题类型, 与热刚度矩阵的问题类型保持一致
	 */
	public void applyHeatSource(ScalarFunc sourceFunc, double[] param, BVPType type) {
		fs.assembleSource(new ScalarFunc[]{sourceFunc}, null, type, RHS);
	}


	public static void main(String[] args) {

		/*
		Mesh2DT3Generator mg = new Mesh2DT3Generator();
		Mesh2DT3 mesh2D = mg.square2D(50, 50, T3Type.Left);

		GeneralCoefFunc ConstFunc = (xyz, relabel, param) -> 1.0e0;
		GeneralCoefFunc SourceFunc = (xyz, relabel, param) -> 2 * xyz[0] * (1 - xyz[0]) + 2 * xyz[1] * (1 - xyz[1]);
		Possion2DT31 p2DT31 = new Possion2DT31(mesh2D);
		p2DT31.DirichletProblem(BVPType.COMMON, ConstFunc, null, SourceFunc, null, null, null, null, "MyResult.dat");
		*/
		
		/*
		Mesh2DT3Generator mg = new Mesh2DT3Generator();
		Mesh2DT3 mesh2D = mg.square2D(40, 40, T3Type.Left);
		GeneralCoefFunc ConstFunc = (xyz, relabel, param) -> 1.0e0;
		GeneralCoefFunc SourceFunc = (xyz, relabel, param) -> 2 * Math.sin(2 * Math.PI * xyz[0]) * Math.sin(2 * Math.PI * xyz[1]);
		int[] dirichletLabel = new int[] {1, 3, 4};
		int[] neumannLabel = new int[] {2};
		GeneralCoefFunc fluxFunc = (xyz, relabel, param) -> xyz[1];
		Possion2DT31 p2DT31 = new Possion2DT31(mesh2D);
		p2DT31.GeneralPossionProblem(BVPType.COMMON, ConstFunc, null,
				                     SourceFunc, null, 
				                     dirichletLabel, null, null, 
				                     neumannLabel, fluxFunc, null, 
				                     null, null, null, null, null, null, 
				                     "Demon.dat");
		*/
		
		/*
		int N = 8;
		Mesh2DT3Generator  mgt3 = new Mesh2DT3Generator();
		Mesh2DT3 fineMesh = mgt3.square2D(256, 256, T3Type.Left);
		RegionFunc regionFineFunc = (xy, param) -> {
            double xCell = xy[0] * N - (int)(xy[0] * N);
            double yCell = xy[1] * N - (int)(xy[1] * N);
            return (Math.abs(xCell - 0.5) < 0.25 && Math.abs(yCell - 0.5) < 0.25)? 1.0:-1.0;
        };
        fineMesh.setDomainLabel(regionFineFunc, null, 10);
        
        GeneralCoefFunc cf = (xyz, label, param) -> label == 10? 0.01: 1;
        GeneralCoefFunc sourceFunc = (xy, llabel, param) -> 10.0;
        Possion2DT31 p2dt31 = new Possion2DT31(fineMesh);
        RegionFunc bdFunc = (xy, param) -> 0.0;
        p2dt31.DirichletProblem(BVPType.COMMON, cf, null, sourceFunc, null,
        		                new int[] {1, 2, 3, 4}, bdFunc, null, "MyNewFineSolution.dat");
//        double[] fineSol = p2dt31.exportSolution();
 */

		ScalarFunc a = (xy, label, param) -> {
//			double x = xy[0];
//			double y = xy[1];
//			if(x < 0) {
//				if(y < 0) {
//					return 1.0;
//				}else {
//					return 0.001;
//				}
//			}else {
//				if (y > 0) {
//					return 1.0;
//				}
//				else {

//					return 0.001;
//				}
//			}
			return 1.0;
		};
		ScalarFunc source = (xy, label, param) -> 1.0;
		Mesh2T3 mesh = new Mesh2T3().square2D(100, 100, T3Type.Left);
//		mesh.scale(2.0, 2.0);
//		mesh.translate(-1.0, -1.0);
		Possion2DT31 p2DT31 = new Possion2DT31(mesh);

		p2DT31.DirichletProblem(BVPType.COMMON, a, null, source, null, null, null, null, "interfaces.dat");



/*
		Mesh2DT3Generator mg = new Mesh2DT3Generator();
		Mesh2DT3 mesh2D = mg.square2D(20, 20, T3Type.Left);

		ScalarCoefFunc ConstFunc = (xyz, relabel, param) -> 1.0e0;
		ScalarCoefFunc SourceFunc = (xyz, relabel, param) -> 1.0e0;
		Possion2DT31 p2DT31 = new Possion2DT31(mesh2D);
		p2DT31.DirichletProblem(BVPType.COMMON, ConstFunc, null, SourceFunc, null, null, null, null, "MyResult.dat");
*/

		/*
		Mesh2T3Generator mg = new Mesh2T3Generator();
		Mesh2T3 mesh2D = mg.square2D(800, 800, T3Type.Left);
		//mesh2D.CMKRelabelling();

		MatrixFunc ConstFunc = (xyz, relabel, param) -> new double[][] {new double[]{1.0, 0.0}, new double[]{0.0, 1.0}};
		ScalarFunc SourceFunc = (xyz, relabel, param) -> 1.0e0;
		Possion2DT31 p2DT31 = new Possion2DT31(mesh2D);
		p2DT31.DirichletProblem(BVPType.CONSTITUTE_MATRIX_COMMON, ConstFunc, null, SourceFunc, null, null, null, null, "MyResult.dat");
		*/
	}

	/**
	 * Possion方程Dirichlet问题
	 * <p>
	 * &part;<sub>i</sub>(a<sub>ij</sub>&part;<sub>j</sub>u) = f &#915;
	 * 
	 * @param kxxFunc
	 * @param paramKxx
	 * @param sourceFunc
	 * @param paramSource
	 * @param tp
	 * @param DirichletBoundaryLabel
	 * @param fixedFunc
	 * @param paramFixed
	 * @param reportFile
	 */
	public void DirichletProblem(BVPType tp, 
			ScalarFunc kxxFunc, double[] paramKxx, 
			ScalarFunc sourceFunc, double[] paramSource, 
			int[] DirichletBoundaryLabel, ScalarFunc fixedFunc, double[] paramFixed,
			String reportFile) {
		fs.assembleStiff(new ScalarFunc[] { kxxFunc }, paramKxx, tp, A);
		if (sourceFunc != null) {
			fs.assembleSource(new ScalarFunc[] { sourceFunc }, paramSource, tp, RHS);
		}
//		System.out.println(A);
//		System.exit(1);
		fs.applyBC_MBN(A, Direct.All, DirichletBoundaryLabel);

		if (fixedFunc == null) {
			fs.applyBC_MBN(RHS, Direct.All, 0.0, DirichletBoundaryLabel);
		} else {
			fs.applyBC_MBN(RHS, Direct.All, new ScalarFunc[] { fixedFunc }, paramFixed, DirichletBoundaryLabel);
		}
		A.save("SRMatrixFile.txt");
		solveAndComputeGradientAndExport(reportFile);
	}

	public void DirichletProblem(BVPType tp, double[][] vecKxx, 
			ScalarFunc sourceFunc, double[] paramSource, 
			int[] DirichletBoundaryLabel, ScalarFunc fixedFunc, double[] paramFixed,
			String reportFile) {
		fs.assembleStiff(vecKxx, tp, A);
		if (sourceFunc != null) {
			fs.assembleSource(new ScalarFunc[] { sourceFunc }, paramSource, tp, RHS);
		}
		fs.applyBC_MBN(A, Direct.All, DirichletBoundaryLabel);

		if (fixedFunc == null) {
			fs.applyBC_MBN(RHS, Direct.All, 0.0, DirichletBoundaryLabel);
		} else {
			fs.applyBC_MBN(RHS, Direct.All, new ScalarFunc[] { fixedFunc }, paramFixed, DirichletBoundaryLabel);
		}
		solveAndComputeGradientAndExport(reportFile);
	}

	public void DirichletProblem(BVPType tp, double[][] vecKxx, double[][] vecSource, 
			int[] DirichletBoundaryLabel, ScalarFunc fixedFunc, double[] paramFixed,
			String reportFile) {
		fs.assembleStiff(vecKxx, tp, A);
		if (vecSource != null) {
			fs.assembleSource(vecSource, tp, RHS);
		}
		fs.applyBC_MBN(A, Direct.All, DirichletBoundaryLabel);

		if (fixedFunc == null) {
			fs.applyBC_MBN(RHS, Direct.All, 0.0, DirichletBoundaryLabel);
		} else {
			fs.applyBC_MBN(RHS, Direct.All, new ScalarFunc[] { fixedFunc }, paramFixed, DirichletBoundaryLabel);
		}

		solveAndComputeGradientAndExport(reportFile);
	}
	
	
	public void DirichletProblem(BVPType tp, 
			MatrixFunc kxxFunc, double[][][] paramKxx,
			ScalarFunc sourceFunc, double[] paramSource, 
			int[] DirichletBoundaryLabel, ScalarFunc fixedFunc, double[] paramFixed,
			String reportFile) {
		fs.assembleStiff(kxxFunc, paramKxx, tp, A);
		if (sourceFunc != null) {
			fs.assembleSource(new ScalarFunc[] { sourceFunc }, paramSource, tp, RHS);
		}
		fs.applyBC_MBN(A, Direct.All, DirichletBoundaryLabel);

		if (fixedFunc == null) {
			fs.applyBC_MBN(RHS, Direct.All, 0.0, DirichletBoundaryLabel);
		} else {
			fs.applyBC_MBN(RHS, Direct.All, new ScalarFunc[] { fixedFunc }, paramFixed, DirichletBoundaryLabel);
		}
		//A.save("SRMatrixFile.txt");
		A.save("MatFile.txt");
		MVO.saveArray("RHSFile.txt", RHS);
		solveAndComputeGradientAndExport(reportFile);
	}
	/**
	 * 
	 * @param tp
	 * @param kxx
	 * @param source
	 * @param DirichletBoundaryLabel
	 * @param fixedBV
	 * @param reportFile
	 */
	public void DirichletProblem(BVPType tp, double[] kxx, double source, int[] DirichletBoundaryLabel,
			double fixedBV, String reportFile) {
		fs.assembleStiff(kxx, tp, A);
		fs.assembleSource(new double[] { source }, tp, RHS);
		fs.applyBC_MBN(A, Direct.All, DirichletBoundaryLabel);
		fs.applyBC_MBN(RHS, Direct.All, fixedBV, DirichletBoundaryLabel);
		solveAndComputeGradientAndExport(reportFile);
	}

	public void GeneralPossionProblem(BVPType tp, ScalarFunc kxxFunc, double[] paramKxx, ScalarFunc sourceFunc,
			double[] paramSource, int[] DirichletBoundaryLabel, ScalarFunc fixedFunc, double[] paramFixed,
			int[] NeumannBoundaryLabel, ScalarFunc fluxFunc, double[] paramFlux, int[] RobinBoundaryLabel,
			ScalarFunc robinFluxFunc, double[] paramRobinFlux, 
			ScalarFunc convectFunc, double[] paramConvectFunc, BVPType robinType, String reportFile) {
		
		fs.assembleStiff(new ScalarFunc[] { kxxFunc }, paramKxx, tp, A);
		
		if (sourceFunc != null) {
			fs.assembleSource(new ScalarFunc[] { sourceFunc }, paramSource, tp, RHS);
		}
		if (NeumannBoundaryLabel != null) {
			fs.assembleFlux(new ScalarFunc[] { fluxFunc }, paramFlux, tp, RHS, NeumannBoundaryLabel);
		}
		if (RobinBoundaryLabel != null) {
			fs.assembleBoundaryMass(new ScalarFunc[] { convectFunc }, paramConvectFunc, robinType, A,
					RobinBoundaryLabel);
			fs.assembleFlux(new ScalarFunc[] { robinFluxFunc }, paramRobinFlux, robinType, RHS, RobinBoundaryLabel);
		}

		if (DirichletBoundaryLabel != null) {
			fs.applyBC_MBN(A, Direct.All, DirichletBoundaryLabel);
			if (fixedFunc != null) {
				fs.applyBC_MBN(RHS, Direct.X, new ScalarFunc[] { fixedFunc }, paramFixed, DirichletBoundaryLabel);
			} else {
				fs.applyBC_MBN(RHS, Direct.X, 0.0, DirichletBoundaryLabel);
			}
		}
		solveAndComputeGradientAndExport(reportFile);
	}

	public void GeneralPossionProblem(BVPType tp, double kxx, double source, int[] DirichletBoundaryLabel,
			double fixedBV, int[] NeumannBoundaryLabel, double fluxValue, int[] RobinBoundaryLabel,
			double robinFluxValue, double convectValue, String reportFile) {
		fs.assembleStiff(new double[] { kxx }, tp, A);
		fs.assembleSource(new double[] { source }, tp, RHS);
		if (NeumannBoundaryLabel != null) {
			fs.assembleFlux(new double[] { fluxValue }, tp, RHS, NeumannBoundaryLabel);
		}
		if (RobinBoundaryLabel != null) {
			fs.assembleBoundaryMass(new double[] { convectValue }, tp, A, RobinBoundaryLabel);
			fs.assembleFlux(new double[] { robinFluxValue }, tp, RHS, RobinBoundaryLabel);
		}

		fs.applyBC_MBN(A, Direct.All, DirichletBoundaryLabel);
		fs.applyBC_MBN(RHS, Direct.X, fixedBV, DirichletBoundaryLabel);
		
		solveAndComputeGradientAndExport(reportFile);
	}
	

	private void solveAndComputeGradientAndExport(String reportFile) {
//		NewIterSSolver solver = new NewIterSSolver(A, false);
//		solver.PCGSSSOR(RHS, x, 1.50, 1);
		DirectSSolver solver = new DirectSSolver(A);
		long begin = System.nanoTime();
		solver.Factorize_Cholesky(1);
		x = solver.Solve_Cholesky(RHS, 0);
		long end = System.nanoTime();
		System.out.printf("min: %20.10f\t max: %20.10f\n", MVO.min(x), MVO.max(x));
		System.out.println("Solve time: " + (end - begin) / 1.0e9 + " s");
//		A.sort();
//		System.out.println(A);
		
//		Arrays.fill(x, 0.0);
//		DirectSSolver dsolver = new DirectSSolver(A, fs.getnBand());
//		dsolver.PLUSolve(RHS, x, 1);
		
//		DirectSSolver dsolver = new DirectSSolver(A);
//		dsolver.PPLUSolve(RHS, x, 1);
		
		fs.computeGradient(x, gradient);
		System.out.println("Solution: ");
		System.out.println("x0 min:" + MVO.min(x) + "\t" + "max: " + MVO.max(x));
		System.out.println("Solution Gradient: ");
		System.out.println("Dx min:" + MVO.min(gradient[0]) + "\t" + "max: " + MVO.max(gradient[0]));
		System.out.println("Dy min:" + MVO.min(gradient[1]) + "\t" + "max: " + MVO.max(gradient[1]));
		double[][] mat = new double[][] { x, gradient[0], gradient[1] };
		if (reportFile != null) {
			mesh.toTecplot(reportFile, mat);
		}
	}
	
	
	public double[] exportSolution() {
		return x.clone();
	}
	
	public double[][] exportGradient() {
		double[][] gd = new double[2][];
		gd[0] = gradient[0].clone();
		gd[1] = gradient[1].clone();
		return gd;
	}
	
	Mesh2T3 mesh;
	FES2T31 fs;
	SRMatrix A;
	double[] RHS;
	double[] x;
	double[][] gradient;

}
